{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# %pdb on\n",
    "import json\n",
    "import re\n",
    "import numpy as np\n",
    "import copy\n",
    "from tqdm import tqdm_notebook \n",
    "from pytorch_pretrained_bert.tokenization import whitespace_tokenize, BasicTokenizer, BertTokenizer\n",
    "from redis import StrictRedis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "db = StrictRedis(host='localhost', port=6379, db=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finish Reading! len =  90447\n"
     ]
    }
   ],
   "source": [
    "with open('./hotpot_train_v1.1.json', 'r') as fin:\n",
    "    train_set = json.load(fin)\n",
    "print('Finish Reading! len = ', len(train_set))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with open('./hotpot_train_v1.1_small.json', 'w') as fout:\n",
    "    json.dump(train_set[:5000], fout)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "GENERAL_WD = ['is', 'are', 'am', 'was', 'were', 'have', 'has', 'had', 'can', 'could', \n",
    "              'shall', 'will', 'should', 'would', 'do', 'does', 'did', 'may', 'might', 'must', 'ought', 'need', 'dare']\n",
    "GENERAL_WD += [x.capitalize() for x in GENERAL_WD]\n",
    "GENERAL_WD = re.compile(' |'.join(GENERAL_WD))\n",
    "\n",
    "def judge_question_type(q : str, G = GENERAL_WD) -> int:\n",
    "    if q.find(' or ') >= 0:\n",
    "        return 2 \n",
    "    elif G.match(q):\n",
    "        return 1\n",
    "    else:\n",
    "        return 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Miami Gardens is a suburban city located in north-central Miami-Dade County, Florida.\n"
     ]
    }
   ],
   "source": [
    "# print(judge_question_type('Who has a longer middle name, Alice Walker or Michael Herr?'))\n",
    "print(db.lrange('Miami Gardens, Florida', 0, -1)[0].decode())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "BERT_MODEL = 'bert-base-uncased'\n",
    "# tokenizer = BertTokenizer.from_pretrained(BERT_MODEL, do_lower_case=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# print(tokenizer.tokenize('Brian \\\"Boosh\\\" Boucher (pronounced \\\"Boo-shay\\\") (born January 2, 1977) is a retired American professional ice hockey goaltender, who played 13 seasons in the National Hockey League (NHL) for the Philadelphia Flyers, Phoenix Coyotes, Calgary Flames, Chicago Blackhawks, Columbus Blue Jackets, San Jose Sharks, and Carolina Hurricanes.'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[('Miami Gardens, Florida', 'Miami Gardens,', 64, 78), ('Hard Rock Stadium', 'Hard Rock Stadium', 0, 17)]\n"
     ]
    }
   ],
   "source": [
    "from hotpot_evaluate_v1 import normalize_answer, f1_score\n",
    "from fuzzywuzzy import fuzz, process as fuzzy_process\n",
    "\n",
    "def fuzzy_retrive(entity, pool):\n",
    "    if len(pool) > 100:\n",
    "        # fullwiki, exact match\n",
    "        # TODO: test ``entity (annotation)'' and find the most like one\n",
    "        if pool.get(entity):\n",
    "            return entity\n",
    "        else:\n",
    "            return None\n",
    "    else:\n",
    "        # distractor mode or use link in original wiki, no need to consider ``entity (annotation)''\n",
    "        pool = pool if isinstance(pool, list) else pool.keys()\n",
    "        f1max, ret = 0, None\n",
    "        for t in pool:\n",
    "            f1, precision, recall = f1_score(entity, t)\n",
    "            if f1 > f1max:\n",
    "                f1max, ret = f1, t\n",
    "        return ret\n",
    "\n",
    "def find_near_matches(w, sentence):\n",
    "    ret = []\n",
    "    max_ratio = 0\n",
    "    t = 0\n",
    "    for word in sentence.split():\n",
    "        while sentence[t] != word[0]:\n",
    "            t += 1\n",
    "        score = (fuzz.ratio(w, word) + fuzz.partial_ratio(w, word)) / 2\n",
    "        if score > max_ratio:\n",
    "            max_ratio = score\n",
    "            ret = [(t, t + len(word))]\n",
    "        elif score == max_ratio:\n",
    "            ret.append((t, t + len(word)))\n",
    "        else:\n",
    "            pass\n",
    "        t += len(word)\n",
    "    return ret if max_ratio > 85 else []     \n",
    "\n",
    "def dp(a, b): # a source, b long text\n",
    "    f, start = np.zeros((len(a), len(b))), np.zeros((len(a), len(b)), dtype = np.int)\n",
    "    for j in range(len(b)):\n",
    "        f[0, j] = int(a[0] != b[j])\n",
    "        if j > 0 and b[j - 1].isalnum():\n",
    "            f[0, j] += 10\n",
    "        start[0, j] = j\n",
    "    for i in range(1, len(a)):        \n",
    "        for j in range(len(b)):\n",
    "            # (0, i-1) + del(i) ~ (start[j], j)\n",
    "            f[i, j] = f[i - 1, j] + 1\n",
    "            start[i, j] = start[i - 1, j]\n",
    "            if j == 0:\n",
    "                continue\n",
    "            if f[i, j] > f[i - 1, j - 1] + int(a[i] != b[j]):\n",
    "                f[i, j] = f[i - 1, j - 1] + int(a[i] != b[j])\n",
    "                start[i, j] = start[i-1, j - 1]\n",
    "\n",
    "            if f[i, j] > f[i, j - 1] + 0.5:\n",
    "                f[i, j] = f[i, j - 1] + 0.5\n",
    "                start[i, j] = start[i, j - 1]\n",
    "#     print(f[len(a) - 1])\n",
    "    r = np.argmin(f[len(a) - 1])\n",
    "    ret = [start[len(a) - 1, r], r + 1]\n",
    "#     print(b[ret[0]:ret[1]])\n",
    "    score = f[len(a) - 1, r] / len(a)\n",
    "    return (ret, score)\n",
    "\n",
    "def fuzzy_find(entities, sentence):\n",
    "    ret = []\n",
    "    for entity in entities:\n",
    "        item = re.sub(r'\\(.*\\)$', '', entity).strip()\n",
    "        r, score = dp(item, sentence)\n",
    "        if score < 0.5:\n",
    "            matched = sentence[r[0]: r[1]].lower()\n",
    "            final_word = item.split()[-1]\n",
    "            # from end\n",
    "            retry = False\n",
    "            while fuzz.partial_ratio(final_word.lower(), matched) < 80:\n",
    "                retry = True\n",
    "                end = len(item) - len(final_word)\n",
    "                while end > 0 and item[end - 1].isspace():\n",
    "                    end -= 1\n",
    "                if end == 0:\n",
    "                    retry = False\n",
    "                    score = 1\n",
    "                    break\n",
    "                item = item[:end]\n",
    "                final_word = item.split()[-1]\n",
    "            if retry:\n",
    "#                 print(entity + ' ### ' + sentence[r[0]: r[1]] + ' ### ' + item)\n",
    "                r, score = dp(item, sentence)\n",
    "                score += 0.1\n",
    "\n",
    "            if score >= 0.5:\n",
    "#                 print(entity + ' ### ' + sentence[r[0]: r[1]] + ' ### ' + item)\n",
    "                continue\n",
    "            del final_word\n",
    "            # from start\n",
    "            retry = False\n",
    "            first_word = item.split()[0]\n",
    "            while fuzz.partial_ratio(first_word.lower(), matched) < 80:\n",
    "                retry = True\n",
    "                start = len(first_word)\n",
    "                while start < len(item) and item[start].isspace():\n",
    "                    start += 1\n",
    "                if start == len(item):\n",
    "                    retry = False\n",
    "                    score = 1\n",
    "                    break\n",
    "                item = item[start:]\n",
    "                first_word = item.split()[0]\n",
    "            if retry:\n",
    "#                 print(entity + ' ### ' + sentence[r[0]: r[1]] + ' ### ' + item)\n",
    "                r, score = dp(item, sentence)\n",
    "                score = max(score, 1 - ((r[1] - r[0]) / len(entity)))\n",
    "                score += 0.1\n",
    "#             if score > 0.5:\n",
    "#                 print(entity + ' ### ' + sentence[r[0]: r[1]] + ' ### ' + item)\n",
    "            if score < 0.5:\n",
    "                if item.isdigit() and sentence[r[0]: r[1]] != item:\n",
    "                    continue\n",
    "                ret.append((entity, sentence[r[0]: r[1]], int(r[0]), int(r[1]), score))\n",
    "    non_intersection = []\n",
    "    for i in range(len(ret)):\n",
    "        ok = True\n",
    "        for j in range(len(ret)):\n",
    "            if j != i:\n",
    "                if not (ret[i][2] >= ret[j][3] or ret[j][2] >= ret[i][3]) and ret[j][4] < ret[i][4]:\n",
    "                    ok = False\n",
    "                    break\n",
    "                if ret[i][4] > 0.2 and ret[j][4] < 0.1 and not ret[i][1][0].isupper() and len(ret[i][1].split()) <= 3:\n",
    "                    ok = False\n",
    "                    print(ret[i])\n",
    "                    break\n",
    "        if ok:\n",
    "            non_intersection.append(ret[i][:4])\n",
    "    return non_intersection\n",
    "\n",
    "# print(dp('Skiffle', 'Die Rh\\u00f6ner S\\u00e4uw\\u00e4ntzt are a Skif, dm , fle-Bluesband from Eichenzell-L\\u00fctter in Hessen, Germany.'))\n",
    "# def fuzzy_find(entities, sentence):\n",
    "#     items = fuzzy_process.extract(sentence, entities, scorer=fuzz.partial_token_set_ratio)\n",
    "#     items = [x for x, y in items if y > 85]\n",
    "#     items_matched = []\n",
    "#     for item in items:\n",
    "#         positions = []\n",
    "#         for w in re.split('[\\s,.?!]', item):\n",
    "#             r = find_near_matches(w, sentence)\n",
    "#             if len(r) > 0:\n",
    "#                 # assume by default sorted by starts\n",
    "#                 positions.append(r)\n",
    "#         # To find an interval, which length is minimized\n",
    "#         print(item, positions)\n",
    "#         assert len(positions) > 0\n",
    "#         min_len, s_min, e_min = len(sentence), -1, -1\n",
    "#         while s_min < 0:\n",
    "#             if len(positions) == 1:\n",
    "#                 s_min, e_min = positions[0][0]\n",
    "#                 break\n",
    "#             for s0, e0 in positions[0]:\n",
    "#                 for s_1, e_1 in positions[-1]:\n",
    "#                     if s_1 <= e0:\n",
    "#                         continue\n",
    "#                     if e_1 - s0 >= min_len:\n",
    "#                         break\n",
    "#                     ok = True\n",
    "# #                 last = e0\n",
    "# #                 for k in range(1, len(positions) - 1):\n",
    "# #                     ok = False\n",
    "# #                     for s_k, e_k in positions[k]:\n",
    "# #                         if last < s_k and e_k < s_1:\n",
    "# #                             last = e_k\n",
    "# #                             ok = True\n",
    "# #                             break\n",
    "# #                     if not ok:\n",
    "# #                         break\n",
    "#                     if ok:\n",
    "#                         min_len, s_min, e_min = e_1 - s0, s0, e_1\n",
    "#             if min_len > 2 * len(item): # invalid, too long\n",
    "#                 positions.pop()\n",
    "#                 s_min, e_min = -1, -1\n",
    "#         items_matched.append(sentence[s_min: e_min])\n",
    "#     return list(zip(items, items_matched))   \n",
    "print(list(fuzzy_find(['Miami Gardens, Florida', 'WSCV', 'Hard Rock Stadium'], r\"Hard Rock Stadium is a multipurpose football stadium located in Miami Gardens, a city north of Miami. It is the home stadium of the Miami Dolphins of the National Football League (NFL).\")))\n",
    "# print(fuzzy_find([\"19 Kids and Counting\", \"nine girls and 10 boys\"], r\" A spin-off show of \\\"19 kids ande counting\\\", it features the Duggar family: Jill Dillard, Jessa Seewald, sixteen of their seventeen siblings, and parents Jim Bob and Michelle Duggar.\"))\n",
    "# print(fuzzy_retrive('Joshua Aaron Charles', ['Jawahar Navodaya Vidyalaya Kanpur', 'Dead Poets Society', 'Josh Charles', 'Aaron1', 'josh charles']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3fc39bd79ca6442485ffc1d68bdeebff",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=90447), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('2000–01 Utah Jazz season', '2000–01 NBA season', 4, 22, 0.375)\n",
      "('Operation Paperclip', 'operation', 123, 132, 0.2111111111111111)\n",
      "('John of Bohemia', 'of Bohemia', 59, 69, 0.43333333333333335)\n",
      "('Tyndall Air Force Base', 'air force base', 6, 20, 0.4636363636363636)\n",
      "('Hurricane Ivan', 'hurricane', 16, 25, 0.2111111111111111)\n",
      "5a7b23ca554299042af8f703\n",
      "('Clark County, Nevada', 'county in Nevada', 5, 21, 0.3142857142857143)\n",
      "('Early flying machines', 'flying machines', 32, 47, 0.3857142857142857)\n",
      "('Us Weekly', 'weekly', 47, 53, 0.43333333333333335)\n",
      "('the quiet Beatle', 'the Beatle', 55, 65, 0.375)\n",
      "('Texas country music', 'country music', 40, 53, 0.4157894736842105)\n",
      "('The Bends', 'the', 89, 92, 0.43333333333333335)\n",
      "('Football in Munich', 'footballin', 6, 16, 0.28181818181818186)\n",
      "('the physical universe', 'the Vienna Univers', 185, 203, 0.42857142857142855)\n",
      "('Spectre (2015 film)', 'actre', 24, 29, 0.42857142857142855)\n",
      "('County Cork', 'county', 36, 42, 0.26666666666666666)\n",
      "('Robins Air Force Base', 'air force base', 28, 42, 0.43333333333333335)\n",
      "('Battle of Lund', 'battle', 22, 28, 0.26666666666666666)\n",
      "('Ánima Estudios', 'animation studio', 5, 21, 0.35714285714285715)\n",
      "('Jefferson County, Kentucky', 'county in Kentucky', 57, 75, 0.4076923076923077)\n",
      "('Montclair State University', 'is the university', 12, 29, 0.4461538461538461)\n",
      "('School of the Air', 'school', 35, 41, 0.26666666666666666)\n",
      "('Division of Barton', 'division', 13, 21, 0.225)\n",
      "('YG Entertainment', 'entertainment', 44, 57, 0.2875)\n",
      "('The Informant', 'the', 14, 17, 0.43333333333333335)\n",
      "('Battle of Reading (871)', 'battle', 6, 12, 0.26666666666666666)\n",
      "('Battle of Verdun', 'battle of', 77, 86, 0.2111111111111111)\n",
      "('Road Trip', 'road trip', 133, 142, 0.2222222222222222)\n",
      "('Olathe School District', 'school district', 35, 50, 0.4181818181818182)\n",
      "('the Member of Parliament (MP)', 'the UK Parliament', 88, 105, 0.375)\n",
      "('STMicroelectronics', 'electronics', 56, 67, 0.4444444444444444)\n",
      "('six', 'si', 63, 65, 0.3333333333333333)\n",
      "('York University', 'university', 148, 158, 0.43333333333333335)\n",
      "('Mount Hotham', 'mount', 48, 53, 0.30000000000000004)\n",
      "('Fyre Festival', 'festival', 104, 112, 0.48461538461538456)\n",
      "('Formula 5000', 'formula', 32, 39, 0.24285714285714285)\n",
      "('society', 'soci', 93, 97, 0.42857142857142855)\n",
      "('Telecom Italia', 'telecom', 35, 42, 0.24285714285714285)\n",
      "('Absurdist fiction', 'absurdist', 31, 40, 0.2111111111111111)\n",
      "('Torpedo boat', 'torpedo', 34, 41, 0.24285714285714285)\n",
      "('Convention on Biological Diversity', 'of biological diversity', 49, 72, 0.4235294117647058)\n",
      "('Fyre Festival', 'festival', 65, 73, 0.48461538461538456)\n",
      "('2017 NHL Expansion Draft', 'expansion draft', 171, 186, 0.475)\n",
      "('The Economist', 'the', 80, 83, 0.43333333333333335)\n",
      "('NATO phonetic alphabet', 'phonetic alphabet', 118, 135, 0.32727272727272727)\n",
      "('Stock car racing', 'stock car', 39, 48, 0.2111111111111111)\n",
      "('Fahd of Saudi Arabia', 'of Saudi Arabia', 133, 148, 0.35)\n",
      "('Battle of Baxter Springs', 'of Baxter Springs', 94, 111, 0.3916666666666666)\n",
      "('The River Wild', 'the Univer', 70, 80, 0.37777777777777777)\n",
      "('EU Business School', 'business school', 48, 63, 0.2666666666666666)\n",
      "('Guitar Player', 'guitar', 45, 51, 0.26666666666666666)\n",
      "('The Freak', 'the', 4, 7, 0.43333333333333335)\n",
      "5abed6d45542990832d3a0ef\n",
      "('Hunter', 'inter', 11, 16, 0.3333333333333333)\n",
      "('Interstate 70 in Kansas', 'interstate that in', 22, 40, 0.35)\n",
      "('9 to 12 April 1917', '9 April 1917', 96, 108, 0.3333333333333333)\n",
      "('located in Tirana, Albania.', 'in Tirana, Albania', 26, 44, 0.43333333333333335)\n",
      "('1970s', '197', 171, 174, 0.4)\n",
      "('Light pollution in Hong Kong', 'light pollution con', 44, 63, 0.2388888888888889)\n",
      "('Swiss electronic', 'electronic', 27, 37, 0.475)\n",
      "('Rumble in the Bronx', 'in the Bronx', 42, 54, 0.46842105263157896)\n",
      "('LG Electronics', 'electronics', 77, 88, 0.3142857142857143)\n",
      "('House of Habsburg', 'houses of', 105, 114, 0.2875)\n",
      "5ab6b2fb5542995eadef0060\n",
      "('Neo soul', 'neo-soul', 6, 14, 0.25)\n",
      "('Neo soul', 'neo-soul', 56, 64, 0.25)\n",
      "('Division of New England', 'division', 19, 27, 0.225)\n",
      "('IPhone 4S', 'phone', 68, 73, 0.43333333333333335)\n",
      "('Journal of Fluid Mechanics', 'of fluid mechanics', 121, 139, 0.4076923076923077)\n",
      "('River Thames', 'river', 143, 148, 0.30000000000000004)\n",
      "('Ardis Publishing', 'publishing', 24, 34, 0.475)\n",
      "('The Amazing Spider-Man 2', 'the \"Amazing Spiderman', 29, 51, 0.2590909090909091)\n",
      "('British', 'artis', 34, 39, 0.42857142857142855)\n",
      "('Roc Nation singles discography', 'singles discography', 4, 23, 0.4666666666666667)\n",
      "('I Am Australian', 'the Australian', 51, 65, 0.26666666666666666)\n",
      "('Crime Wave', 'crime', 170, 175, 0.30000000000000004)\n",
      "('boxer', 'box', 32, 35, 0.4)\n",
      "('Acro dance', 'acrobatic dance', 92, 107, 0.35)\n",
      "('Trance music', 'trance', 43, 49, 0.26666666666666666)\n",
      "('La Rioja Province, Argentina', 'province in Argentina', 40, 61, 0.35)\n",
      "('Sukhoi Design Bureau', 'design bureau', 107, 120, 0.44999999999999996)\n",
      "('aviation', 'aviato', 69, 75, 0.25)\n",
      "('The Theatre', 'theatre', 31, 38, 0.45454545454545453)\n",
      "('President Ronald Reagan', 'president', 83, 92, 0.2111111111111111)\n",
      "('Horror Stories 2', 'horror', 76, 82, 0.26666666666666666)\n",
      "('Rome Protocols', 'protocols', 56, 65, 0.4571428571428571)\n",
      "('Mount Abu', 'mount', 16, 21, 0.30000000000000004)\n",
      "('1980s', '198', 64, 67, 0.4)\n",
      "('Elena of Montenegro', 'of Montenegro', 46, 59, 0.4157894736842105)\n",
      "('Anna University', 'university', 29, 39, 0.43333333333333335)\n",
      "('The Wave (2015 film)', 'the', 170, 173, 0.43333333333333335)\n",
      "('Emory University', 'university', 8, 18, 0.475)\n",
      "('forward', 'award', 52, 57, 0.42857142857142855)\n",
      "5ae0e2df5542990adbacf6b1\n",
      "('1981 NFL season', '1981 season', 254, 265, 0.26666666666666666)\n",
      "('Interstate 81 in Virginia', 'interstate ', 16, 27, 0.475)\n",
      "('the Pacific Ocean', 'northern Pacific Ocean', 321, 343, 0.2647058823529412)\n",
      "('Executive Order 13769', 'executive order', 29, 44, 0.23333333333333334)\n",
      "('The Nightingale', 'nightingale', 197, 208, 0.3666666666666667)\n",
      "('News Corporation', 'corporation', 280, 291, 0.4125)\n",
      "('Vikings', 'viking', 34, 40, 0.2857142857142857)\n",
      "('Urban Dictionary', 'dictionary', 45, 55, 0.475)\n",
      "('Selena', 'relea', 15, 20, 0.3333333333333333)\n",
      "('Starz', 'star', 15, 19, 0.4)\n",
      "('forward', 'award', 186, 191, 0.42857142857142855)\n",
      "('Electronic musical instrument', 'electronic instrument', 57, 78, 0.3103448275862069)\n",
      "('Hurricane Irene', 'hurricane', 57, 66, 0.2111111111111111)\n",
      "('The Louvin Brothers', 'the', 51, 54, 0.43333333333333335)\n",
      "('New York Stock Exchange', 'stock exchange', 58, 72, 0.4913043478260869)\n",
      "('June deportation', 'deportation', 47, 58, 0.4125)\n",
      "('Winter X Games XV', 'winter games', 15, 27, 0.3857142857142857)\n",
      "('Fictional country', 'fictional', 30, 39, 0.2111111111111111)\n",
      "('Montana University System', 'university system', 18, 35, 0.41999999999999993)\n",
      "('University of Copenhagen', 'univeristy', 60, 70, 0.35)\n",
      "('2017 NHL Expansion Draft', 'expansion draft', 279, 294, 0.475)\n",
      "('quetzal and other feathers', 'quetzal feather headdres', 199, 223, 0.36538461538461536)\n",
      "('Film industry', 'industry', 5, 13, 0.48461538461538456)\n",
      "('Hurricane Sandy', 'hurricane', 10, 19, 0.2111111111111111)\n",
      "('Relapse 2', 'release', 55, 62, 0.4571428571428572)\n",
      "('1850s', '185', 32, 35, 0.4)\n",
      "('1780s', '178', 21, 24, 0.4)\n",
      "('Bank of England', 'bank of', 71, 78, 0.24285714285714285)\n",
      "('Wholesaling', 'wholesal', 229, 237, 0.36363636363636365)\n",
      "('Recluse spider', 'recluse', 49, 56, 0.24285714285714285)\n",
      "('M62 motorway', 'motorway', 29, 37, 0.43333333333333335)\n",
      "('Trinity Church Cemetery', 'church cemetery', 43, 58, 0.4478260869565217)\n",
      "('Serenity Role Playing Game', 'role playing game', 36, 53, 0.4461538461538461)\n",
      "('Starz', 'star', 15, 19, 0.4)\n",
      "('Battle of the Coral Sea', 'battle betwe', 210, 222, 0.48461538461538467)\n",
      "('Liberal Party of Australia', 'party in Australia', 174, 192, 0.4076923076923077)\n",
      "('National Health Service', 'health service', 31, 45, 0.4913043478260869)\n",
      "('Trolls (film)', 'dolls', 58, 63, 0.3333333333333333)\n",
      "('Gulf of Maine', 'gulf of', 11, 18, 0.24285714285714285)\n",
      "('Book', 'book', 44, 48, 0.25)\n",
      "('PeteStrumentals', 'instrumental', 122, 134, 0.4)\n",
      "('Touro Law Center', 'law center', 49, 59, 0.475)\n",
      "('Coles Supermarkets', 'supermarket', 34, 45, 0.4888888888888888)\n",
      "('Demographics of Canada', 'demographic', 56, 67, 0.26666666666666666)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('jazz tenor saxophonist', 'tenor saxophonist', 106, 123, 0.32727272727272727)\n",
      "('Olathe School District', 'school district', 48, 63, 0.4181818181818182)\n",
      "('Australia national soccer team', 'national soccer team', 51, 71, 0.43333333333333335)\n",
      "('Dean of Lincoln', 'of Lincoln', 43, 53, 0.43333333333333335)\n",
      "('Special Boat Service', 'special', 34, 41, 0.24285714285714285)\n",
      "('Computer security', 'computer', 67, 75, 0.225)\n",
      "5a8d6138554299585d9e37c7\n",
      "('The Guard (2011 film)', 'the', 41, 44, 0.43333333333333335)\n",
      "('Franco Malerba', 'doFranco Malerba', 5, 21, 0.21428571428571427)\n",
      "('Please (Pet Shop Boys album)', 'release', 48, 55, 0.4166666666666667)\n",
      "('Eraring Power Station', 'power station', 13, 26, 0.4809523809523809)\n",
      "('M1 carbine', 'carbine', 52, 59, 0.4)\n",
      "('Journal of Applied Physics', 'journal', 5, 12, 0.24285714285714285)\n",
      "('NXP Semiconductors', 'semiconductor', 59, 72, 0.37777777777777777)\n",
      "('12 million', 'million', 114, 121, 0.4)\n",
      "('About Time (2013 film)', 'about', 57, 62, 0.30000000000000004)\n",
      "('Mobile app', 'mobile', 60, 66, 0.26666666666666666)\n",
      "('Pizzagate conspiracy theory', 'conspiracy theory', 49, 66, 0.4703703703703703)\n",
      "('in Macau and southern China', 'in Macau, China', 106, 121, 0.48148148148148145)\n",
      "('Game Center', 'game', 52, 56, 0.35)\n",
      "('Ohio Turnpike', 'turnpike', 46, 54, 0.48461538461538456)\n",
      "('Shen Yun Performing Arts', 'performing-arts', 40, 55, 0.475)\n",
      "('Kansas Legislature', 'legislature', 35, 46, 0.4888888888888888)\n",
      "('The Everly Brothers', 'the brothers', 22, 34, 0.47368421052631576)\n",
      "('Reims Cathedral', 'catherdral', 45, 55, 0.43333333333333335)\n",
      "('Niki destinations', 'destinations', 41, 53, 0.39411764705882346)\n",
      "('STMicroelectronics', 'electronics', 54, 65, 0.4444444444444444)\n",
      "('MacBook', 'iBook', 62, 67, 0.42857142857142855)\n",
      "('Yucca Mountain nuclear waste repository', 'nuclear waste repository', 6, 30, 0.48461538461538456)\n",
      "('Slasher film', 'slasher horror film', 40, 59, 0.375)\n",
      "('The Crown (TV series)', 'the', 142, 145, 0.43333333333333335)\n",
      "5ab740165542992aa3b8c7fa\n",
      "('Characters of Overwatch', 'character of', 38, 50, 0.25384615384615383)\n",
      "('seventh', 'seven', 173, 178, 0.2857142857142857)\n",
      "('Urban Dictionary', 'dictionary', 35, 45, 0.475)\n",
      "('Patreon', 'patron', 399, 405, 0.2857142857142857)\n",
      "('Man Booker International Prize', 'international prize', 49, 68, 0.4666666666666667)\n",
      "(\"Indiana's 9th congressional district\", 'congressional district', 5, 27, 0.4888888888888888)\n",
      "('Mash ingredients', 'ingredient', 42, 52, 0.475)\n",
      "('Teletubbies', 'andTeletubbies', 44, 58, 0.3181818181818182)\n",
      "('Disney Interactive', 'interactive', 57, 68, 0.4888888888888888)\n",
      "('Opera seria', 'opera', 27, 32, 0.30000000000000004)\n",
      "('1970s', '197', 93, 96, 0.4)\n",
      "('George and Ira Gershwin', 'and the Gershwin', 41, 57, 0.4043478260869565)\n",
      "('Shoegazing', 'shoegaz', 65, 72, 0.4)\n",
      "('Shoegazing', 'shoegaz', 65, 72, 0.4)\n",
      "('Internet fraud', 'internet', 18, 26, 0.225)\n",
      "('Washington metropolitan area', 'metropolitan area', 18, 35, 0.4928571428571429)\n",
      "('seven', 'serve', 22, 27, 0.3)\n",
      "('ten', 'te', 91, 93, 0.3333333333333333)\n",
      "('Chinese Journal of Physics', 'journal of physics', 46, 64, 0.4076923076923077)\n",
      "('Sony/ATV Music Publishing', 'music publishing', 23, 39, 0.45999999999999996)\n",
      "('Argentina', 'agent in', 90, 98, 0.3888888888888889)\n",
      "('1960s', '196', 15, 18, 0.4)\n",
      "('Blacks', 'black', 110, 115, 0.3333333333333333)\n",
      "('PC Magazine', 'magazine', 37, 45, 0.3727272727272727)\n",
      "('England national football team', 'international football', 55, 77, 0.3666666666666667)\n",
      "('LG Corporation', 'corporation', 14, 25, 0.3142857142857143)\n",
      "('Australia', 'australi', 165, 173, 0.2222222222222222)\n",
      "('the noon sun', 'the sun', 132, 139, 0.4166666666666667)\n",
      "('Tropical Storm Delta (2005)', 'tropical storm', 18, 32, 0.24285714285714285)\n",
      "('Football in Belarus', 'football', 37, 45, 0.225)\n",
      "('LG Electronics', 'electronics', 80, 91, 0.3142857142857143)\n",
      "('Song 2', 'song', 25, 29, 0.35)\n",
      "('relax', 'relea', 143, 148, 0.3)\n",
      "('University of Saskatchewan', 'university, of', 4, 18, 0.2153846153846154)\n",
      "5ab2f812554299545a2cfaee\n",
      "('novelist', 'novel', 49, 54, 0.375)\n",
      "('Hamm Building', 'building', 5, 13, 0.48461538461538456)\n",
      "('Internment', 'interne', 134, 141, 0.4)\n",
      "('1930s', '193', 10, 13, 0.4)\n",
      "('National personification', 'personification', 76, 91, 0.475)\n",
      "('DirecTV', 'direc', 100, 105, 0.42857142857142855)\n",
      "('Opera seria', 'opera', 15, 20, 0.30000000000000004)\n",
      "('Emory University', 'university', 31, 41, 0.475)\n",
      "('The Sea', 'the', 51, 54, 0.43333333333333335)\n",
      "('Subprefecture', 'prefecture', 212, 222, 0.3076923076923077)\n",
      "('Nelson County, Kentucky', 'county in Kentucky', 5, 23, 0.31739130434782603)\n",
      "('A divorce', 'divorce', 36, 43, 0.3222222222222222)\n",
      "('Thompson submachine gun', 'sub-machine gun', 34, 49, 0.4478260869565217)\n",
      "('Disney Interactive', 'interactive', 14, 25, 0.4888888888888888)\n",
      "('States of Germany', 'state of', 64, 72, 0.3222222222222222)\n",
      "('Shooting of Eddie Hutch Snr', 'shoothing of', 55, 67, 0.23636363636363636)\n",
      "(\"New York's 20th congressional district\", 'congressional distric t', 4, 27, 0.49473684210526314)\n",
      "('Super Bowl XLV', 'super bowl', 51, 61, 0.30000000000000004)\n",
      "('Comune', 'commune', 30, 37, 0.25)\n",
      "('Pachinko', 'machin', 43, 49, 0.375)\n",
      "('Orange County, Florida', 'county in Florida', 89, 106, 0.32727272727272727)\n",
      "('Coors Brewing Company', 'brewing company', 20, 35, 0.3857142857142857)\n",
      "('York University', 'university', 60, 70, 0.43333333333333335)\n",
      "('The More', 'the', 13, 16, 0.43333333333333335)\n"
     ]
    }
   ],
   "source": [
    "# construct cognitive graph in training data    \n",
    "def find_fact_content(bundle, title, sen_num):\n",
    "    for x in bundle['context']:\n",
    "        if x[0] == title:\n",
    "            return x[1][sen_num]\n",
    "test = copy.deepcopy(train_set)\n",
    "for bundle in tqdm_notebook(test):\n",
    "    entities = set([title for title, sen_num in bundle['supporting_facts']])\n",
    "    bundle['Q_edge'] = fuzzy_find(entities, bundle['question'])\n",
    "    for fact in bundle['supporting_facts']:\n",
    "        try:\n",
    "            title, sen_num = fact\n",
    "            pool = set()\n",
    "            for i in range(sen_num + 1):\n",
    "                name = 'edges:###{}###{}'.format(i, title)\n",
    "                tmp = set([x.decode().split('###')[0] for x in db.lrange(name, 0, -1)])\n",
    "                pool |= tmp\n",
    "            pool &= entities\n",
    "            pool.add(bundle['answer'])\n",
    "            pool.discard(title)\n",
    "            r = fuzzy_find(pool, find_fact_content(bundle, title, sen_num))\n",
    "            fact.append(r)\n",
    "        except IndexError as e:\n",
    "            print(bundle['_id'])\n",
    "with open('./hotpot_train_v1.1_refined.json', 'w') as fout:\n",
    "    json.dump(test, fout)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
